package com.example.audio_test.utils

import android.content.Context
import android.media.MediaCodec
import android.media.MediaCodecInfo
import android.media.MediaFormat
import android.os.Environment
import kotlinx.coroutines.GlobalScope
import kotlinx.coroutines.delay
import kotlinx.coroutines.launch
import java.io.BufferedOutputStream
import java.io.File
import java.io.FileOutputStream
import java.io.IOException
import java.lang.Exception
import java.nio.ByteBuffer
import java.util.concurrent.ArrayBlockingQueue

/**
 *H264编码类
 *
 * @author : Kun
 * @date : 2021/6/21
 * @time : 17:16
 */
class H264Encoder(var width: Int, var height: Int, var framerate: Int, var context: Context) {

    companion object {
        val TIMEOUT_USEC = 12000L
    }

    private var mediaCodec: MediaCodec? = null

    var isRunning = false

    private var configByte: ByteArray? = null

    private var outputStream: BufferedOutputStream? = null

    /**
     * 缓存队列，缓存每一帧的图像
     */
    private var yuv420Queue = ArrayBlockingQueue<ByteArray>(10)

    init {
        val mediaFormat = MediaFormat.createVideoFormat("video/avc", width, height).apply {
            setInteger(
                MediaFormat.KEY_COLOR_FORMAT,
                MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420Flexible
            )
            setInteger(MediaFormat.KEY_BIT_RATE, width * height * 5)
            setInteger(MediaFormat.KEY_FRAME_RATE, framerate)
            setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, 1)
        }

        mediaCodec = MediaCodec.createEncoderByType("video/avc")
        mediaCodec!!.configure(mediaFormat, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE)
        mediaCodec!!.start()

        createFile()
    }

    /**
     * 创建文件
     */
    private fun createFile() {
        val file = File(context.getExternalFilesDir(Environment.DIRECTORY_MOVIES), "test.h264")
        if (file.exists()) {
            file.delete()
        }
        outputStream = BufferedOutputStream(FileOutputStream(file))
    }

    /**
     * 添加数据
     */
    fun putData(bufferData: ByteArray) {
        if (yuv420Queue.size >= 10) {
            yuv420Queue.poll()
        }
        yuv420Queue.add(bufferData)
    }

    /**
     * 开始编码
     */
    fun startEncoder() {
        GlobalScope.launch {
            isRunning = true
            var input: ByteArray? = null
            var pts = 0L
            var generateIndex = 0L

            while (isRunning) {
                if (yuv420Queue.size > 0) {
                    input = yuv420Queue.poll()
                    nv21ToNV12(input, width, height)
                }
                if (input != null) {
                    /**
                     * 延迟 TIME_US 等待拿到空的 input buffer下标，单位为 us
                     * -1 表示一直等待，直到拿到数据，0 表示立即返回
                     */
                    val inputBufferIndex = mediaCodec!!.dequeueInputBuffer(-1)
                    if (inputBufferIndex >= 0) {
                        //获取一个可用的空的 inputBuffer
                        val inputBuffer = mediaCodec!!.getInputBuffer(inputBufferIndex)
                        inputBuffer!!.clear()
                        inputBuffer.put(input)
                        //处理输入视频帧
                        mediaCodec!!.queueInputBuffer(
                            inputBufferIndex,
                            0,
                            input.size,
                            System.currentTimeMillis(),
                            0
                        )
                    }

                    //每个缓冲区元数据包括指定相关编解码器（输出）缓冲区中有效数据范围的偏移量和大小。
                    val bufferInfo = MediaCodec.BufferInfo()
                    var outputBufferIndex =
                        mediaCodec!!.dequeueOutputBuffer(bufferInfo, TIMEOUT_USEC)
                    //循环获取输出数据
                    while (outputBufferIndex >= 0) {
                        //获取处理好的outputBuffer
                        val outputBuffer = mediaCodec!!.getOutputBuffer(outputBufferIndex)
                        val outData = ByteArray(bufferInfo.size)
                        //输出数据写到数组里
                        outputBuffer!!.get(outData)

                        if (bufferInfo.flags == MediaCodec.BUFFER_FLAG_CODEC_CONFIG) {
                            //这表明标记为这样的缓冲区包含编解码器初始化/编解码器特定数据而不是媒体数据
                            //视频数据头文件？
                            configByte = ByteArray(outData.size)
                            configByte = outData
                        } else if (bufferInfo.flags == MediaCodec.BUFFER_FLAG_KEY_FRAME) {
                            //这表示标记为这样的（编码）缓冲区包含关键帧的数据。
                            val keyFrame = ByteArray(configByte!!.size + outData.size)
                            System.arraycopy(configByte, 0, keyFrame, 0, configByte!!.size)
                            System.arraycopy(outData, 0, keyFrame, configByte!!.size, outData.size)
                            try {
                                outputStream!!.write(keyFrame, 0, keyFrame.size)
                            } catch (e: Exception) {
                                e.printStackTrace()
                            }
                        } else {
                            try {
                                outputStream!!.write(outData, 0, outData.size)
                            } catch (e: Exception) {
                                e.printStackTrace()
                            }
                        }
                        mediaCodec!!.releaseOutputBuffer(outputBufferIndex, false)
                        outputBufferIndex = mediaCodec!!.dequeueOutputBuffer(
                            bufferInfo,
                            TIMEOUT_USEC
                        )
                    }
                } else {
                    delay(500)
                }
            }
            try {
                //停止解码器释放资源
                mediaCodec!!.stop()
                mediaCodec!!.release()
            } catch (e: Exception) {
                e.printStackTrace()
            }
            try {
                outputStream!!.flush()
                outputStream!!.close()
            } catch (e: IOException) {
                e.printStackTrace()
            }
        }
    }

    /**
     * 停止编码
     */
    fun stopEncoder(){
        isRunning=false
    }

    /**
     * NV21转NV21
     *
     *
     *  YUV，分为三个分量，“Y”表示明亮度（Luminance或Luma），也就是灰度值；
     * 而“U”和“V” 表示的则是色度（Chrominance或Chroma），作用是描述影像色彩及饱和度，用于指定像素的颜色。
     * YUV码流的存储格式其实与其采样的方式密切相关，主流的采样方式有三种，YUV4:4:4，YUV4:2:2，YUV4:2:0，
     * YUV 4:4:4采样，每一个Y对应一组UV分量。
     *YUV 4:2:2采样，每两个Y共用一组UV分量。
     *YUV 4:2:0采样，每四个Y共用一组UV分量。
     *
     * NV12存储格式: YYYYYYYY UVUV     =>YUV420SP
     * NV21存储格式: YYYYYYYY VUVU     =>YUV420SP
     *
     * 更多了解  https://www.cnblogs.com/azraelly/archive/2013/01/01/2841269.html
     *
     * 这里采用的是420
     * 转换时只需交换UV
     *
     */
    private fun nv21ToNV12(input: ByteArray?, width: Int, height: Int) {
        if (input == null) {
            return
        }

        /*
        *  这里采用的是4:2:0
        *   所需内存大小 Y=width*height  U = Y/4 V = Y/4 (共计 width*height*3/2)
         */
        val frameSize = width * height
        var temp: Byte
        for (i in 2 until frameSize / 2 step 2) {
            //前部分为Y不用管，后部分交换UV;UV交替互換
            //把V赋值给temp
            temp = input[frameSize - 1 + i]
            //把改位置上的V变成U
            input[frameSize - 1 + i] = input[frameSize - 1 + i - 1]
            //U变成V
            input[frameSize - 1 + i - 1] = temp
        }
    }
}